What did I do this week?
This week, I focused on resolving compatibility issues caused by the recent upgrade to jax==0.6.2
. The update removed internal access to attributes like jaxlib.xla_extension.PjitFunction
and pxla._xla.PjitFunction
, which previously allowed us to identify JAX-compiled functions. This change caused test failures due to broken type checks.
To resolve the issue, I updated the result.py
module by replacing the old type-checking logic with a safer alternative:
= type(jax.jit(lambda x: x)) _jitfun_type
This avoids reliance on private APIs and ensures future-proof compatibility across JAX versions.
I re-verified the failing tests with the updated `qutip-jax` installed, and **all tests now pass**.
Pull Request Progress
After completing the fix, I addressed all pending review comments, removed unnecessary print and debug statements, and wrapped up the final discussion point. I then submitted the pull request, which successfully passed all tests.
QTRL Logging Cleanup in QOC
Alongside the JAX fix, I continued working on the QTRL-to-QOC migration effort by cleaning up legacy logging code in QOC’s optimizer modules.
🔧 Specifically, I:
- Removed QTRL-specific logging usage from
_grape.py
and_crab.py
. - Eliminated:
import qutip_qtrl.logging_utils
logger = logging.get_logger()
- All
logger.debug(...)
calls
I ensured that surrounding logic tied to self._qtrl.stats
, self._qtrl.iter_summary
, and self._qtrl.dump
remained untouched, so functionality was preserved even as unused logging code was removed.
After making these changes, I ran the full test suite and confirmed that all tests passed.
I’m excited to be moving closer to a cleaner, more modular, and JAX-compatible QOC codebase that’s easier to maintain and extend for future optimizers.